#!/usr/bin/env python3
"""
SDXL LoRA训练参数配置节点
专门针对SDXL模型优化的训练参数配置
基于SDXL的特性和最佳实践设计
"""

import os
import psutil
import torch
import folder_paths

# 导入统一定义的参数类型
from .lora_trainer_utils.constants import SDXL_TRAINING_PARAMS, SDXL_TRAINING_PARAMS_CATEGORY

class SDXLTrainingParams:
    """SDXL LoRA训练参数配置节点"""
    
    aux_id = "comfyui_lora_train"
    
    @classmethod
    def INPUT_TYPES(cls):
        # 获取系统信息
        total_memory = psutil.virtual_memory().total / (1024 ** 3)  # GB
        cpu_count = psutil.cpu_count(logical=False)  # 物理CPU核心数
        
        # 获取CUDA信息
        cuda_available = torch.cuda.is_available()
        if cuda_available:
            gpu_name = torch.cuda.get_device_name(0)
            gpu_memory = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)  # GB
        else:
            gpu_name = "未检测到CUDA设备"
            gpu_memory = 0
            
        system_info = f"系统内存: {total_memory:.1f}GB\nCPU核心数: {cpu_count}\nGPU: {gpu_name}\nGPU显存: {gpu_memory:.1f}GB"
        
        return {
            "required": {
                # === 系统信息显示 ===
                "system_info": ("STRING", {
                    "default": system_info,
                    "multiline": True,
                    "readonly": True,
                    "tooltip": "当前系统硬件信息"
                }),
                
                # === 训练阶段选择 ===
                "training_stage": (["第一阶段-基础训练", "第二阶段-精细训练", "第三阶段-高质量训练", "自定义参数"], {
                    "default": "第二阶段-精细训练",
                    "tooltip": "选择训练阶段，不同阶段有不同的参数配置"
                }),
                
                # === 核心训练参数 ===
                "train_batch_size": ("INT", {
                    "default": 1,
                    "min": 1,
                    "max": 4,
                    "step": 1,
                    "tooltip": "训练批次大小，SDXL显存需求较大，建议1-2"
                }),
                
                "learning_rate": ("FLOAT", {
                    "default": 0.0001,
                    "min": 0.00001,
                    "max": 0.001,
                    "step": 0.00001,
                    "tooltip": "学习率，SDXL建议使用较小的学习率"
                }),
                
                "network_dim": ("INT", {
                    "default": 64,
                    "min": 16,
                    "max": 256,
                    "step": 8,
                    "tooltip": "网络维度，SDXL需要更大的维度以获得更好效果"
                }),
                
                "network_alpha": ("INT", {
                    "default": 64,
                    "min": 8,
                    "max": 256,
                    "step": 4,
                    "tooltip": "网络Alpha，通常设为network_dim的0.5-1倍"
                }),
                
                "max_train_steps": ("INT", {
                    "default": 1500,
                    "min": 200,
                    "max": 10000,
                    "step": 100,
                    "tooltip": "最大训练步数，SDXL需要更多步数"
                }),
                
                # === 训练控制参数 ===
                "save_every_n_steps": ("INT", {
                    "default": 300,
                    "min": 100,
                    "max": 1000,
                    "step": 50,
                    "tooltip": "每N步保存一次模型"
                }),
                
                "gradient_accumulation_steps": ("INT", {
                    "default": 2,
                    "min": 1,
                    "max": 8,
                    "step": 1,
                    "tooltip": "梯度累积步数，SDXL建议使用2-4"
                }),
                
                "resolution": ("INT", {
                    "default": 1024,
                    "min": 512,
                    "max": 1536,
                    "step": 64,
                    "tooltip": "训练分辨率，SDXL推荐1024或更高"
                }),
                
                # === SDXL专用参数 ===
                "clip_skip": ("INT", {
                    "default": 2,
                    "min": 1,
                    "max": 4,
                    "step": 1,
                    "tooltip": "CLIP跳过层数，SDXL推荐2"
                }),
                
                "face_crop_aug_range": ("FLOAT", {
                    "default": 0.5,
                    "min": 0.1,
                    "max": 1.0,
                    "step": 0.1,
                    "tooltip": "人脸裁剪增强范围，提高人物特征学习"
                }),
                
                "caption_dropout_rate": ("FLOAT", {
                    "default": 0.1,
                    "min": 0.0,
                    "max": 0.5,
                    "step": 0.05,
                    "tooltip": "标签丢弃率，提高模型泛化能力"
                }),
                
                # === 数据增强参数 ===
                "flip_aug": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "启用水平翻转增强"
                }),
                
                "color_aug": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "启用颜色增强（SDXL训练通常关闭）"
                }),
                
                "random_crop": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "启用随机裁剪（SDXL训练通常关闭）"
                }),
                
                # === 优化器参数 ===
                "optimizer_type": (["AdamW8bit", "AdamW", "Lion", "SGDNesterov"], {
                    "default": "AdamW8bit",
                    "tooltip": "优化器类型，8bit版本节省显存"
                }),
                
                "lr_scheduler": (["cosine", "linear", "constant", "constant_with_warmup"], {
                    "default": "cosine",
                    "tooltip": "学习率调度器，cosine通常效果最好"
                }),
                
                "lr_warmup_steps": ("INT", {
                    "default": 100,
                    "min": 0,
                    "max": 1000,
                    "step": 10,
                    "tooltip": "学习率预热步数，SDXL建议100-200"
                }),
                
                # === SDXL特定参数 ===
                "network_train_unet_only": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "仅训练UNet，不训练文本编码器"
                }),
                
                "network_train_text_encoder_only": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "仅训练文本编码器，不训练UNet"
                }),
                
                "prior_loss_weight": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.0,
                    "max": 10.0,
                    "step": 0.1,
                    "tooltip": "先验损失权重"
                }),
                
                "min_snr_gamma": ("FLOAT", {
                    "default": 5.0,
                    "min": 0.0,
                    "max": 20.0,
                    "step": 0.5,
                    "tooltip": "最小信噪比gamma值"
                }),
                
                "scale_weight_norms": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.0,
                    "max": 10.0,
                    "step": 0.1,
                    "tooltip": "权重范数缩放"
                }),
                
                # === SDXL高级参数 ===
                "cache_text_encoder_outputs": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "缓存文本编码器输出，SDXL推荐启用"
                }),
                
                "cache_text_encoder_outputs_to_disk": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "将文本编码器输出缓存到磁盘"
                }),
                
                "enable_bucket": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "启用分辨率桶，提高训练效率"
                }),
                
                "bucket_reso_steps": ("INT", {
                    "default": 64,
                    "min": 32,
                    "max": 128,
                    "step": 8,
                    "tooltip": "分辨率桶步长"
                }),
                
                "bucket_no_upscale": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "分辨率桶不放大"
                }),
                
                "max_token_length": ("INT", {
                    "default": 225,
                    "min": 75,
                    "max": 225,
                    "step": 1,
                    "tooltip": "最大token长度，SDXL最大225"
                }),
            },
            "optional": {
                # === 可选高级参数 ===
                "mixed_precision": (["fp16", "bf16", "no"], {
                    "default": "fp16",
                    "tooltip": "混合精度训练类型，SDXL推荐fp16"
                }),
                
                "gradient_checkpointing": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "梯度检查点，节省显存"
                }),
                
                "cache_latents": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "缓存潜在表示，提高训练速度"
                }),
                
                "persistent_data_loader_workers": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "持久化数据加载器工作进程"
                }),
                
                "save_precision": (["fp16", "bf16", "float"], {
                    "default": "fp16",
                    "tooltip": "模型保存精度"
                }),
                
                "max_train_epochs": ("INT", {
                    "default": 0,
                    "min": 0,
                    "max": 100,
                    "step": 1,
                    "tooltip": "最大训练轮数，0表示使用步数控制"
                }),
                
                "seed": ("INT", {
                    "default": 42,
                    "min": 0,
                    "max": 2147483647,
                    "step": 1,
                    "tooltip": "随机种子"
                }),
            }
        }
    
    RETURN_TYPES = ("DICT",)
    RETURN_NAMES = ("SDXL训练参数",)
    FUNCTION = "generate_sdxl_params"
    CATEGORY = SDXL_TRAINING_PARAMS_CATEGORY
    
    def generate_sdxl_params(self, system_info, training_stage, train_batch_size, learning_rate,
                           network_dim, network_alpha, max_train_steps, save_every_n_steps,
                           gradient_accumulation_steps, resolution, clip_skip, face_crop_aug_range,
                           caption_dropout_rate, flip_aug, color_aug, random_crop, optimizer_type,
                           lr_scheduler, lr_warmup_steps, network_train_unet_only,
                           network_train_text_encoder_only, prior_loss_weight, min_snr_gamma,
                           scale_weight_norms, cache_text_encoder_outputs, cache_text_encoder_outputs_to_disk,
                           enable_bucket, bucket_reso_steps, bucket_no_upscale, max_token_length,
                           mixed_precision="fp16", gradient_checkpointing=True, cache_latents=True,
                           persistent_data_loader_workers=True, save_precision="fp16",
                           max_train_epochs=0, seed=42):
        """生成SDXL LoRA训练参数"""
        
        # 根据训练阶段自动调整参数
        if training_stage == "第一阶段-基础训练":
            # 快速基础训练
            auto_params = {
                "learning_rate": 0.0005,
                "network_dim": 32,
                "network_alpha": 32,
                "max_train_steps": 800,
                "save_every_n_steps": 200,
                "face_crop_aug_range": 0.3,
                "caption_dropout_rate": 0.2,
                "lr_warmup_steps": 50,
            }
        elif training_stage == "第二阶段-精细训练":
            # 标准精细训练
            auto_params = {
                "learning_rate": 0.0001,
                "network_dim": 64,
                "network_alpha": 64,
                "max_train_steps": 1500,
                "save_every_n_steps": 300,
                "face_crop_aug_range": 0.5,
                "caption_dropout_rate": 0.1,
                "lr_warmup_steps": 100,
            }
        elif training_stage == "第三阶段-高质量训练":
            # 高质量训练
            auto_params = {
                "learning_rate": 0.00012,
                "network_dim": 128,
                "network_alpha": 128,
                "max_train_steps": 3000,
                "save_every_n_steps": 300,
                "face_crop_aug_range": 0.7,
                "caption_dropout_rate": 0.05,
                "lr_warmup_steps": 200,
            }
        else:
            # 自定义参数，使用用户输入
            auto_params = {}
        
        # 合并自动参数和用户参数
        sdxl_params = {
            # 核心训练参数
            "train_batch_size": train_batch_size,
            "learning_rate": auto_params.get("learning_rate", learning_rate),
            "network_dim": auto_params.get("network_dim", network_dim),
            "network_alpha": auto_params.get("network_alpha", network_alpha),
            "max_train_steps": auto_params.get("max_train_steps", max_train_steps),
            "save_every_n_steps": auto_params.get("save_every_n_steps", save_every_n_steps),
            
            # 训练控制参数
            "gradient_accumulation_steps": gradient_accumulation_steps,
            "resolution": resolution,
            
            # SDXL专用参数
            "clip_skip": clip_skip,
            "face_crop_aug_range": auto_params.get("face_crop_aug_range", face_crop_aug_range),
            "caption_dropout_rate": auto_params.get("caption_dropout_rate", caption_dropout_rate),
            
            # 数据增强参数
            "flip_aug": flip_aug,
            "color_aug": color_aug,
            "random_crop": random_crop,
            
            # 优化器参数
            "optimizer_type": optimizer_type,
            "lr_scheduler": lr_scheduler,
            "lr_warmup_steps": auto_params.get("lr_warmup_steps", lr_warmup_steps),
            
            # 高级参数
            "network_train_unet_only": network_train_unet_only,
            "network_train_text_encoder_only": network_train_text_encoder_only,
            "prior_loss_weight": prior_loss_weight,
            "min_snr_gamma": min_snr_gamma,
            "scale_weight_norms": scale_weight_norms,
            
            # SDXL特定参数
            "cache_text_encoder_outputs": cache_text_encoder_outputs,
            "cache_text_encoder_outputs_to_disk": cache_text_encoder_outputs_to_disk,
            "enable_bucket": enable_bucket,
            "bucket_reso_steps": bucket_reso_steps,
            "bucket_no_upscale": bucket_no_upscale,
            "max_token_length": max_token_length,
            
            # 可选参数
            "mixed_precision": mixed_precision,
            "gradient_checkpointing": gradient_checkpointing,
            "cache_latents": cache_latents,
            "persistent_data_loader_workers": persistent_data_loader_workers,
            "save_precision": save_precision,
            "max_train_epochs": max_train_epochs,
            "seed": seed,
            
            # SDXL固定参数
            "v2": False,  # SDXL不是v2
            "v_parameterization": False,  # SDXL不使用v参数化
            "sdxl": True,  # 启用SDXL模式
            
            # SDXL特定参数
            "disable_mmap_load_safetensors": False,  # SDXL特定参数
        }
        
        return (sdxl_params,)


class SimpleSDXLTrainingParams:
    """SDXL LoRA基础训练参数配置节点"""
    
    aux_id = "comfyui_lora_train"
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                # === 核心训练参数 ===
                "train_batch_size": ("INT", {
                    "default": 1,
                    "min": 1,
                    "max": 4,
                    "step": 1,
                    "tooltip": "训练批次大小，SDXL显存需求较大，建议1-2"
                }),
                
                "learning_rate": ("FLOAT", {
                    "default": 0.0001,
                    "min": 0.00001,
                    "max": 0.001,
                    "step": 0.00001,
                    "tooltip": "学习率，SDXL建议使用较小的学习率"
                }),
                
                "network_dim": ("INT", {
                    "default": 64,
                    "min": 16,
                    "max": 256,
                    "step": 8,
                    "tooltip": "网络维度，SDXL需要更大的维度以获得更好效果"
                }),
                
                "network_alpha": ("INT", {
                    "default": 64,
                    "min": 8,
                    "max": 256,
                    "step": 4,
                    "tooltip": "网络Alpha，通常设为network_dim的0.5-1倍"
                }),
                
                "max_train_steps": ("INT", {
                    "default": 1500,
                    "min": 200,
                    "max": 10000,
                    "step": 100,
                    "tooltip": "最大训练步数，SDXL需要更多步数"
                }),
                
                "save_every_n_steps": ("INT", {
                    "default": 300,
                    "min": 100,
                    "max": 1000,
                    "step": 50,
                    "tooltip": "每N步保存一次模型"
                }),
                
                # === 训练控制参数 ===
                "gradient_accumulation_steps": ("INT", {
                    "default": 2,
                    "min": 1,
                    "max": 8,
                    "step": 1,
                    "tooltip": "梯度累积步数，SDXL建议使用2-4"
                }),
                
                "resolution": ("INT", {
                    "default": 1024,
                    "min": 512,
                    "max": 1536,
                    "step": 64,
                    "tooltip": "训练分辨率，SDXL推荐1024或更高"
                }),
                
                # === 优化器参数 ===
                "optimizer_type": (["AdamW8bit", "AdamW", "Lion", "SGDNesterov"], {
                    "default": "AdamW8bit",
                    "tooltip": "优化器类型，8bit版本节省显存"
                }),
                
                "lr_scheduler": (["cosine", "linear", "constant", "constant_with_warmup"], {
                    "default": "cosine",
                    "tooltip": "学习率调度器，cosine通常效果最好"
                }),
            },
            "optional": {
                # === 可选参数 ===
                "clip_skip": ("INT", {
                    "default": 2,
                    "min": 1,
                    "max": 4,
                    "step": 1,
                    "tooltip": "CLIP跳过层数，SDXL推荐2"
                }),
                
                "lr_warmup_steps": ("INT", {
                    "default": 100,
                    "min": 0,
                    "max": 1000,
                    "step": 10,
                    "tooltip": "学习率预热步数，SDXL建议100-200"
                }),
            }
        }
    
    RETURN_TYPES = (SDXL_TRAINING_PARAMS,)
    RETURN_NAMES = ("SDXL基础训练参数",)
    FUNCTION = "generate_simple_sdxl_params"
    CATEGORY = SDXL_TRAINING_PARAMS_CATEGORY
    
    def generate_simple_sdxl_params(self, train_batch_size, learning_rate, network_dim, 
                                  network_alpha, max_train_steps, save_every_n_steps,
                                  gradient_accumulation_steps, resolution, optimizer_type,
                                  lr_scheduler, clip_skip=2, lr_warmup_steps=100):
        """生成SDXL LoRA基础训练参数"""
        
        # 构建基础训练参数
        simple_sdxl_params = {
            # 核心训练参数
            "train_batch_size": train_batch_size,
            "learning_rate": learning_rate,
            "network_dim": network_dim,
            "network_alpha": network_alpha,
            "max_train_steps": max_train_steps,
            "save_every_n_steps": save_every_n_steps,
            
            # 训练控制参数
            "gradient_accumulation_steps": gradient_accumulation_steps,
            "resolution": resolution,
            
            # 优化器参数
            "optimizer_type": optimizer_type,
            "lr_scheduler": lr_scheduler,
            "lr_warmup_steps": lr_warmup_steps,
            
            # SDXL专用参数
            "clip_skip": clip_skip,
            
            # 固定参数（SDXL训练推荐值）
            "seed": 42,
            "mixed_precision": "fp16",
            "gradient_checkpointing": True,
            "cache_latents": True,
            "cache_text_encoder_outputs": True,
            "cache_text_encoder_outputs_to_disk": False,
            "enable_bucket": True,
            "bucket_reso_steps": 64,
            "bucket_no_upscale": True,
            "max_token_length": 225,
            "persistent_data_loader_workers": True,
            "save_precision": "fp16",
            "max_train_epochs": 0,
            
            # SDXL训练增强参数
            "face_crop_aug_range": 0.5,
            "caption_dropout_rate": 0.1,
            "flip_aug": True,
            "color_aug": False,
            "random_crop": False,
            
            # 高级参数
            "network_train_unet_only": False,
            "network_train_text_encoder_only": False,
            "prior_loss_weight": 1.0,
            "min_snr_gamma": 5.0,
            "scale_weight_norms": 1.0,
            
            # SDXL固定参数
            "v2": False,  # SDXL不是v2
            "v_parameterization": False,  # SDXL不使用v参数化
            "sdxl": True,  # 启用SDXL模式
        }
        
        return (simple_sdxl_params,)


# 注册节点
NODE_CLASS_MAPPINGS = {
    "SDXLTrainingParams": SDXLTrainingParams,
    "SimpleSDXLTrainingParams": SimpleSDXLTrainingParams
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SDXLTrainingParams": "🎨 SDXL训练参数",
    "SimpleSDXLTrainingParams": "🎨 SDXL基础参数"
} 